#! /usr/bin/env python3
# -*-coding: utf-8-*-

"""
Created on Apr 8 2017

@author: Moonkie

@attention: 
"""


import mxnet as mx


def ConvolutionPoolingFactory(data,kernel,stride,num_filter,pool_type,pkernel,pstride,covname,actname,poolname,pad=(0,0),ppad=(0,0)):
    conv = mx.sym.Convolution(data=data,kernel=kernel,stride=stride,pad=pad,num_filter=num_filter,name=covname,weight=mx.init.Xavier(),dilate=(1,1))
    relu = mx.sym.Activation(data=conv,name=actname,act_type='relu')
    pool = mx.sym.Pooling(data=relu,name=poolname,pool_type=pool_type,kernel=pkernel,stride=pstride,pad=pad)
    return pool


def ConvolutionReluFactory(data,kernel,stride,num_filter,covname,actname,pad=(0,0) ):
    conv = mx.sym.Convolution(data=data,kernel=kernel,stride=stride,pad=pad,num_filter=num_filter,name=covname)
    relu = mx.sym.Activation(data=conv,name=actname,act_type='relu')
    return relu


def FullyConnectedReluFactory(data,n1,fc1,actname):
    fc1 = mx.sym.FullyConnected(data=data,num_hidden=n1,name=fc1,weight=mx.init.Xavier())
    relu = mx.sym.Activation(data=fc1,name=actname,act_type='relu')
    return relu


def FullyConnectedReluFullyConnectedFactory(data,n1,n2,fc1,fc2,actname):
    fc1 = mx.sym.FullyConnected(data=data,num_hidden=n1,name=fc1,weight=mx.init.Xavier())
    relu = mx.sym.Activation(data=fc1,name=actname,act_type='relu')
    fc2 = mx.sym.FullyConnected(data=relu,name=fc2,num_hidden=n2,weight=mx.init.Xavier())
    return fc2


def PnetFactory(data,kernel1,stride1,bias1,pkernel1,pstride1,kernel2,stride2,bias2,kernel3,stride3,bias3,kernel4,stride4,bias4,ppad=(0,0)):
    conv1 = mx.sym.Convolution(data=data,name='conv1',kernel=kernel1,stride=stride1,weight=mx.init.Xavier(),bias=bias1,num_filter=28)
    relu1 = mx.sym.Activation(data=conv1,name='prelu1',act_type='relu')
    pool1 = mx.sym.Pooling(data=relu1,pool_type='max',name='pool1',kernel=pkernel1,stride=pstride1,pad=ppad)

    conv2 = mx.sym.Convolution(data=pool1,name='conv2',kernel=kernel2,stride=stride2,weight=mx.init.Xavier(),bias=bias2)
    relu2 = mx.sym.Activation(data=conv2,name='prelu2',act_type='relu')

    conv3 = mx.sym.Convolution(data=relu2,name='conv3',kernel=kernel3,stride=stride3,weight=mx.init.Xavier(),bias=bias3)
    relu3 = mx.sym.Activation(data=conv3,name='prelu3',act_type='relu')

    conv41 = mx.sym.Convolution(data=relu3,name='conv41',kernel=kernel4,stride=stride4,weight=mx.init.Xavier(),bias=bias4)
    conv42 = mx.sym.Convolution(data=conv41,name='conv42',kernel=kernel4,stride=stride4,weight=mx.init.Xavier(),bias=bias4)

    out = mx.sym.SoftmaxOutput(data=conv42,name='prob1')
    return out


def RnetFactory(data,kernel1,stride1,bias1,pkernel1,pstride1,kernel2,stride2,bias2,kernel3,stride3,bias3,ppad=(0,0)):
    conv1 = mx.sym.Convolution(data=data,name='conv1',kernel=kernel1,stride=stride1,weight=mx.init.Xavier(),bias=bias1,num_filter=28)
    relu1 = mx.sym.Activation(data=conv1,name='prelu1',act_type='relu')
    pool1 = mx.sym.Pooling(data=relu1,pool_type='max',name='pool1',kernel=pkernel1,stride=pstride1,pad=ppad)

    conv2 = mx.sym.Convolution(data=pool1,name='conv2',kernel=kernel2,stride=stride2,weight=mx.init.Xavier(),bias=bias2,num_filter=48)
    relu2 = mx.sym.Activation(data=conv2,name='prelu2',act_type='relu')
    pool2 = mx.sym.Pooling(data=relu2,pool_type='max',name='pool2')

    conv3 = mx.sym.Convolution(data=relu2,name='conv3',kernel=kernel3,stride=stride3,weight=mx.init.Xavier(),bias=bias3,num_filter=64)
    relu3 = mx.sym.Activation(data=conv3,name='prelu3',act_type='relu')

    conv4 = mx.sym.FullyConnected(data=relu3,name='conv4',num_hidden=128,weight=mx.init.Xavier())
    relu4 = mx.sym.Activation(data=conv4,name='prelu4',act_type='relu')

    conv51 = mx.sym.FullyConnected(data=relu4,num_hidden=2,name='conv5-1')
    conv52 = mx.sym.FullyConnected(data=conv51,num_hidden=2,name='conv5-2')

    out = mx.sym.SoftmaxOutput(data=conv52,name='prob1')
    return out


def OnetFactory(data,kernel1,stride1,bias1,pkernel1,pstride1,kernel2,stride2,bias2,kernel3,stride3,bias3,ppad=(0,0)):
    conv1 = mx.sym.Convolution(data=data,name='conv1',kernel=kernel1,stride=stride1,weight=mx.init.Xavier(),bias=bias1,num_filter=32)
    relu1 = mx.sym.Activation(data=conv1,name='prelu1',act_type='relu')
    pool1 = mx.sym.Pooling(data=relu1,pool_type='max',name='pool1',kernel=pkernel1,stride=pstride1,pad=ppad)

    conv2 = mx.sym.Convolution(data=pool1,name='conv2',kernel=kernel2,stride=stride2,weight=mx.init.Xavier(),bias=bias2,num_filter=64)
    relu2 = mx.sym.Activation(data=conv2,name='prelu2',act_type='relu')
    pool2 = mx.sym.Pooling(data=relu2,pool_type='max',name='pool2')

    conv3 = mx.sym.Convolution(data=relu2,name='conv3',kernel=kernel3,stride=stride3,weight=mx.init.Xavier(),bias=bias3,num_filter=64)
    relu3 = mx.sym.Activation(data=conv3,name='prelu3',act_type='relu')
    pool3 = mx.sym.Pooling(data=relu3,name='pool3',pool_type='max',kernel=(3,3),stride=(2,2))

    conv4 = mx.sym.Convolution(data=pool3,name='conv4',num_hidden=128,num_filter=128)
    relu4 = mx.sym.Activation(data=conv4,name='prelu4',act_type='relu')

    conv5 = mx.sym.FullyConnected(data=relu4,num_hidden=256,name='conv5')

    drop5 = mx.sym.Dropout(data=conv5,name='drop5',p=0.25)
    relu5 = mx.sym.FullyConnected(data=conv5,name='prelu5',act_type='relu')

    conv61 = mx.sym.FullyConnected(data=relu5,num_hidden=2,name='conv6-1')
    conv62 = mx.sym.FullyConnected(data=conv61,num_hidden=4,name='conv6-2')
    conv63 = mx.sym.FullyConnected(data=conv62,num_hidden=10,name='conv6-3')
    out = mx.sym.SoftmaxOutput(data=conv63,name='prob1')
    return out


def LnetFactory(data):
    slicer_data = mx.sym.SliceChannel(data=data,axis=1,num_outputs=5,name='slice')

    pool11 = ConvolutionPoolingFactory(data=slicer_data[0],kernel=(3,3),stride=(1,1),num_filter=28,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv1_1',actname='prelu1_1',poolname='pool1_1')
    pool21 = ConvolutionPoolingFactory(data=pool11,kernel=(3,3),stride=(1,1),num_filter=48,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv2_1',actname='prelu2_1',poolname='pool2_1')
    prelu31 = ConvolutionReluFactory(data=pool21,kernel=(2,2),stride=(1,1),num_filter=64,covname='conv3_1',actname='prelu3_1')

    pool12 = ConvolutionPoolingFactory(data=slicer_data[1],kernel=(3,3),stride=(1,1),num_filter=28,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv1_2',actname='prelu1_2',poolname='pool1_2')
    pool22 = ConvolutionPoolingFactory(data=pool12,kernel=(3,3),stride=(1,1),num_filter=48,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv2_2',actname='prelu2_2',poolname='pool2_2')
    prelu32 = ConvolutionReluFactory(data=pool22,kernel=(2,2),stride=(1,1),num_filter=64,covname='conv3_2',actname='prelu3_2')

    pool13 = ConvolutionPoolingFactory(data=slicer_data[2],kernel=(3,3),stride=(1,1),num_filter=28,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv1_3',actname='prelu1_3',poolname='pool1_3')
    pool23 = ConvolutionPoolingFactory(data=pool13,kernel=(3,3),stride=(1,1),num_filter=48,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv2_3',actname='prelu2_3',poolname='pool2_3')
    prelu33 = ConvolutionReluFactory(data=pool23,kernel=(2,2),stride=(1,1),num_filter=64,covname='conv3_3',actname='prelu3_3')

    pool14 = ConvolutionPoolingFactory(data=slicer_data[3],kernel=(3,3),stride=(1,1),num_filter=28,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv1_4',actname='prelu1_4',poolname='pool1_4')
    pool24 = ConvolutionPoolingFactory(data=pool14,kernel=(3,3),stride=(1,1),num_filter=48,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv2_4',actname='prelu2_4',poolname='pool2_4')
    prelu34 = ConvolutionReluFactory(data=pool24,kernel=(2,2),stride=(1,1),num_filter=64,covname='conv3_4',actname='prelu3_4')

    pool15 = ConvolutionPoolingFactory(data=slicer_data[4],kernel=(3,3),stride=(1,1),num_filter=28,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv1_5',actname='prelu1_5',poolname='pool1_5')
    pool25 = ConvolutionPoolingFactory(data=pool15,kernel=(3,3),stride=(1,1),num_filter=48,pool_type='max',pkernel=(3,3),pstride=(2,2),covname='conv2_5',actname='prelu2_5',poolname='pool2_5')
    prelu35 = ConvolutionReluFactory(data=pool25,kernel=(2,2),stride=(1,1),num_filter=64,covname='conv3_5',actname='prelu3_5')

    concat = mx.sym.Concat(data=slicer_data,name='concat')

    fc4 = FullyConnectedReluFactory(concat,n1=256,fc1='fc4',actname='prelu4')

    fc51 = FullyConnectedReluFullyConnectedFactory(data=fc4,n1=64,n2=2,fc1='fc4_1',fc2='fc5_1',actname='prelu4_1')
    fc52 = FullyConnectedReluFullyConnectedFactory(data=fc4,n1=64,n2=2,fc1='fc4_2',fc2='fc5_2',actname='prelu4_2')
    fc53 = FullyConnectedReluFullyConnectedFactory(data=fc4,n1=64,n2=2,fc1='fc4_3',fc2='fc5_3',actname='prelu4_3')
    fc54 = FullyConnectedReluFullyConnectedFactory(data=fc4,n1=64,n2=2,fc1='fc4_4',fc2='fc5_4',actname='prelu4_4')
    fc55 = FullyConnectedReluFullyConnectedFactory(data=fc4,n1=64,n2=2,fc1='fc4_5',fc2='fc5_5',actname='prelu4_5')
    group = mx.sym.Group([fc51,fc52,fc53,fc54,fc55])
    return group